"""
Finite Difference Operators

Implements discrete differential operators for structured grids with appropriate
boundary treatment for L-shaped domains.
"""

"""
    central_diff_x(f::Matrix{Float64}, i::Int, j::Int, h::Float64) -> Float64

Compute central difference ∂f/∂x at grid point (i,j).
"""
function central_diff_x(f::Matrix{Float64}, i::Int, j::Int, h::Float64)
    return (f[i+1, j] - f[i-1, j]) / (2h)
end

"""
    central_diff_y(f::Matrix{Float64}, i::Int, j::Int, h::Float64) -> Float64

Compute central difference ∂f/∂y at grid point (i,j).
"""
function central_diff_y(f::Matrix{Float64}, i::Int, j::Int, h::Float64)
    return (f[i, j+1] - f[i, j-1]) / (2h)
end

"""
    laplacian_2d(f::Matrix{Float64}, i::Int, j::Int, h::Float64) -> Float64

Compute 5-point Laplacian ∇²f at grid point (i,j).
"""  
function laplacian_2d(f::Matrix{Float64}, i::Int, j::Int, h::Float64)
    return (f[i+1, j] + f[i-1, j] + f[i, j+1] + f[i, j-1] - 4*f[i, j]) / (h*h)
end

# ===================================================================
# 3D Finite Difference Operators
# ===================================================================

"""
    central_diff_x_3d(f::Array{Float64,3}, i::Int, j::Int, k::Int, h::Float64) -> Float64

Compute central difference ∂f/∂x at 3D grid point (i,j,k).
"""
function central_diff_x_3d(f::Array{Float64,3}, i::Int, j::Int, k::Int, h::Float64)
    return (f[i+1, j, k] - f[i-1, j, k]) / (2h)
end

"""
    central_diff_y_3d(f::Array{Float64,3}, i::Int, j::Int, k::Int, h::Float64) -> Float64

Compute central difference ∂f/∂y at 3D grid point (i,j,k).
"""
function central_diff_y_3d(f::Array{Float64,3}, i::Int, j::Int, k::Int, h::Float64)
    return (f[i, j+1, k] - f[i, j-1, k]) / (2h)
end

"""
    central_diff_z_3d(f::Array{Float64,3}, i::Int, j::Int, k::Int, h::Float64) -> Float64

Compute central difference ∂f/∂z at 3D grid point (i,j,k).
"""
function central_diff_z_3d(f::Array{Float64,3}, i::Int, j::Int, k::Int, h::Float64)
    return (f[i, j, k+1] - f[i, j, k-1]) / (2h)
end

"""
    laplacian_3d(f::Array{Float64,3}, i::Int, j::Int, k::Int, h::Float64) -> Float64

Compute 7-point Laplacian ∇²f at 3D grid point (i,j,k).
"""
function laplacian_3d(f::Array{Float64,3}, i::Int, j::Int, k::Int, h::Float64)
    return (f[i+1, j, k] + f[i-1, j, k] + 
            f[i, j+1, k] + f[i, j-1, k] + 
            f[i, j, k+1] + f[i, j, k-1] - 6*f[i, j, k]) / (h*h)
end

"""
    curl_3d(ψx::Array{Float64,3}, ψy::Array{Float64,3}, ψz::Array{Float64,3}, 
            i::Int, j::Int, k::Int, h::Float64) -> Tuple{Float64, Float64, Float64}

Compute curl of vector field ψ = (ψx, ψy, ψz) at grid point (i,j,k).
Returns (u, v, w) = ∇ × ψ.

The curl is:
- u = ∂ψz/∂y - ∂ψy/∂z
- v = ∂ψx/∂z - ∂ψz/∂x  
- w = ∂ψy/∂x - ∂ψx/∂y
"""
function curl_3d(ψx::Array{Float64,3}, ψy::Array{Float64,3}, ψz::Array{Float64,3}, 
                 i::Int, j::Int, k::Int, h::Float64)
    
    # Compute partial derivatives
    dψz_dy = central_diff_y_3d(ψz, i, j, k, h)
    dψy_dz = central_diff_z_3d(ψy, i, j, k, h)
    
    dψx_dz = central_diff_z_3d(ψx, i, j, k, h)
    dψz_dx = central_diff_x_3d(ψz, i, j, k, h)
    
    dψy_dx = central_diff_x_3d(ψy, i, j, k, h)
    dψx_dy = central_diff_y_3d(ψx, i, j, k, h)
    
    # Compute curl components
    u = dψz_dy - dψy_dz
    v = dψx_dz - dψz_dx
    w = dψy_dx - dψx_dy
    
    return u, v, w
end

"""
    compute_velocities_from_stream_3d!(u::Array{Float64,3}, v::Array{Float64,3}, w::Array{Float64,3},
                                       ψx::Array{Float64,3}, ψy::Array{Float64,3}, ψz::Array{Float64,3},
                                       domain::LShapeDomain{3}, h::Float64)

Compute velocity field from vector stream function using curl operation.
u = ∇ × ψ in the fluid domain.
"""
function compute_velocities_from_stream_3d!(u::Array{Float64,3}, v::Array{Float64,3}, w::Array{Float64,3},
                                           ψx::Array{Float64,3}, ψy::Array{Float64,3}, ψz::Array{Float64,3},
                                           domain::LShapeDomain{3}, h::Float64)
    nx, ny, nz = size(u)
    
    @inbounds for k in 2:nz-1, j in 2:ny-1, i in 2:nx-1
        if domain.fluid_mask[i, j, k]
            u[i, j, k], v[i, j, k], w[i, j, k] = curl_3d(ψx, ψy, ψz, i, j, k, h)
        else
            u[i, j, k] = 0.0
            v[i, j, k] = 0.0
            w[i, j, k] = 0.0
        end
    end
    
    # Ensure solid regions have zero velocity
    u[.!domain.fluid_mask] .= 0.0
    v[.!domain.fluid_mask] .= 0.0
    w[.!domain.fluid_mask] .= 0.0
end

"""
    upwind_convection_3d(f::Array{Float64,3}, u::Array{Float64,3}, v::Array{Float64,3}, w::Array{Float64,3},
                         i::Int, j::Int, k::Int, h::Float64) -> Float64

Compute upwind discretization of convection term u·∇f at grid point (i,j,k).
"""
function upwind_convection_3d(f::Array{Float64,3}, u::Array{Float64,3}, v::Array{Float64,3}, w::Array{Float64,3},
                             i::Int, j::Int, k::Int, h::Float64)
    
    # x-direction convection
    if u[i, j, k] >= 0.0
        conv_x = u[i, j, k] * (f[i, j, k] - f[i-1, j, k]) / h
    else
        conv_x = u[i, j, k] * (f[i+1, j, k] - f[i, j, k]) / h
    end
    
    # y-direction convection  
    if v[i, j, k] >= 0.0
        conv_y = v[i, j, k] * (f[i, j, k] - f[i, j-1, k]) / h
    else
        conv_y = v[i, j, k] * (f[i, j+1, k] - f[i, j, k]) / h
    end
    
    # z-direction convection
    if w[i, j, k] >= 0.0
        conv_z = w[i, j, k] * (f[i, j, k] - f[i, j, k-1]) / h
    else
        conv_z = w[i, j, k] * (f[i, j, k+1] - f[i, j, k]) / h
    end
    
    return conv_x + conv_y + conv_z
end

"""
    velocity_gradient_3d(u::Array{Float64,3}, v::Array{Float64,3}, w::Array{Float64,3},
                         i::Int, j::Int, k::Int, h::Float64) -> Matrix{Float64}

Compute velocity gradient tensor ∇u at grid point (i,j,k).
Returns a 3×3 matrix with components [∂ui/∂xj].
"""
function velocity_gradient_3d(u::Array{Float64,3}, v::Array{Float64,3}, w::Array{Float64,3},
                             i::Int, j::Int, k::Int, h::Float64)
    
    ∇u = zeros(3, 3)
    
    # ∂u/∂x, ∂u/∂y, ∂u/∂z
    ∇u[1, 1] = central_diff_x_3d(u, i, j, k, h)
    ∇u[1, 2] = central_diff_y_3d(u, i, j, k, h) 
    ∇u[1, 3] = central_diff_z_3d(u, i, j, k, h)
    
    # ∂v/∂x, ∂v/∂y, ∂v/∂z
    ∇u[2, 1] = central_diff_x_3d(v, i, j, k, h)
    ∇u[2, 2] = central_diff_y_3d(v, i, j, k, h)
    ∇u[2, 3] = central_diff_z_3d(v, i, j, k, h)
    
    # ∂w/∂x, ∂w/∂y, ∂w/∂z
    ∇u[3, 1] = central_diff_x_3d(w, i, j, k, h)
    ∇u[3, 2] = central_diff_y_3d(w, i, j, k, h)
    ∇u[3, 3] = central_diff_z_3d(w, i, j, k, h)
    
    return ∇u
end

"""
    vorticity_stretching_3d(ωx::Array{Float64,3}, ωy::Array{Float64,3}, ωz::Array{Float64,3},
                           u::Array{Float64,3}, v::Array{Float64,3}, w::Array{Float64,3},
                           i::Int, j::Int, k::Int, h::Float64) -> Tuple{Float64, Float64, Float64}

Compute vorticity stretching term (ω·∇)u at grid point (i,j,k).
Returns (Sx, Sy, Sz) components of the stretching term.
"""
function vorticity_stretching_3d(ωx::Array{Float64,3}, ωy::Array{Float64,3}, ωz::Array{Float64,3},
                                u::Array{Float64,3}, v::Array{Float64,3}, w::Array{Float64,3},
                                i::Int, j::Int, k::Int, h::Float64)
    
    # Get velocity gradient tensor
    ∇u = velocity_gradient_3d(u, v, w, i, j, k, h)
    
    # Vorticity vector at this point
    ω = [ωx[i, j, k], ωy[i, j, k], ωz[i, j, k]]
    
    # Compute (ω·∇)u = ω · ∇u (matrix-vector product)
    # This gives the stretching term for each velocity component
    Sx = ω[1] * ∇u[1, 1] + ω[2] * ∇u[1, 2] + ω[3] * ∇u[1, 3]  # (ω·∇)u
    Sy = ω[1] * ∇u[2, 1] + ω[2] * ∇u[2, 2] + ω[3] * ∇u[2, 3]  # (ω·∇)v  
    Sz = ω[1] * ∇u[3, 1] + ω[2] * ∇u[3, 2] + ω[3] * ∇u[3, 3]  # (ω·∇)w
    
    return Sx, Sy, Sz
end